# coding=utf-8
import pandas as pd
import numpy as np
from sqlalchemy import create_engine
from sqlalchemy.pool import NullPool
import os

'''
获取数据库连接 postgresql
return 数据库连接
'''
DIR = os.getcwd() + '/hive2/'


def getConnection(host, port, dbname, user, password):
    # conn = pg.connect(host=host, port=port, dbname=dbname, user=user, password=password)
    engine = create_engine('postgresql://' + user + ':' + password + '@' + host + ':' + str(port) + '/' + dbname,
                           poolclass=NullPool)
    return engine.connect()
    # return conn


'''
获取数据库连接 db2
return 数据库连接
'''


def getConnectionDb2(host, port, dbname, user, password):
    # conn = pg.connect(host=host, port=port, dbname=dbname, user=user, password=password)
    engine = create_engine('ibm_db_sa://' + user + ':' + password + '@' + host + ':' + str(port) + '/' + dbname,
                           encoding="utf-8", poolclass=NullPool)
    return engine.connect()
    # return conn


'''
获取数据库连接 sts
return 数据库连接
'''
def getConnectionSTS(host, port, dbname, user, password):
    import jaydebeapi
    import jpype

    if jpype.isJVMStarted() and not jpype.isThreadAttachedToJVM():
        jpype.attachThreadToJVM()
        jpype.java.lang.Thread.currentThread().setContextClassLoader(jpype.java.lang.ClassLoader.getSystemClassLoader())

    DRIVER = 'org.apache.hive.jdbc.HiveDriver'
    JARSPATH = [DIR + 'commons-logging-1.2.jar'
        , DIR + 'guava-18.0.jar'
        , DIR + 'hadoop-common-2.6.0-cdh5.14.0.jar'
        , DIR + 'hive-common-1.1.0-cdh5.14.0.jar'
        , DIR + 'hive-jdbc-1.1.0-cdh5.14.0.jar'
        , DIR + 'hive-metastore-1.1.0-cdh5.14.0.jar'
        , DIR + 'hive-service-1.1.0-cdh5.14.0.jar'
        , DIR + 'hive-shims-0.23-1.1.0-cdh5.14.0.jar'
        , DIR + 'hive-shims-common-1.1.0-cdh5.14.0.jar'
        , DIR + 'httpclient-4.2.5.jar'
        , DIR + 'httpcore-4.2.5.jar'
        , DIR + 'libfb303-0.9.3.jar'
        , DIR + 'libthrift-0.9.3.jar'
        , DIR + 'log4j-1.2.17.jar'
        , DIR + 'slf4j-api-1.7.7.jar'
        , DIR + 'slf4j-log4j12-1.7.7.jar'

    ]
    # JDBC connection string
    URL = "jdbc:hive2://" + host + ":" + str(port) + "/" + dbname + ";auth=ldap;user="+ user +";password="+ password
    # URL = "jdbc:hive2://10.70.248.187:21050/dbname;auth=ldap;user=admin-bgtasoq000;password=admin"
    # Connect to HiveServer2

    conn = jaydebeapi.connect(DRIVER, URL, [user, password], JARSPATH)
    return conn


##ORACLE, MYSQL, SQLSERVER, SQLSERVER2005, DB2, INFORMIX, SYBASE, OTHER, EMPTY
# 'firebird',
# 'mssql',
# 'mysql',
# 'oracle',
# 'postgresql',
# 'sqlite',
# 'sybase',
def getConnectionDb3(host, port, dbname, user, password, dbtype):
    if dbtype == "DB2":
        engine = create_engine('ibm_db_sa://' + user + ':' + password + '@' + host + ':' + str(port) + '/' + dbname,
                               poolclass=NullPool)
        return engine.connect()
    elif dbtype == "MYSQL":
        engine = create_engine('mysql://' + user + ':' + password + '@' + host + ':' + str(port) + '/' + dbname,
                               encoding='utf-8', echo=True)
        return engine.connect()
    elif dbtype == "POSTGRESQL":
        engine = create_engine('postgresql://' + user + ':' + password + '@' + host + ':' + str(port) + '/' + dbname,
                               encoding="utf-8", poolclass=NullPool)
        return engine.connect()
    elif dbtype == "ORACLE":
        engine = create_engine('oracle://' + user + ':' + password + '@' + host + ':' + str(port) + '/' + dbname,
                               encoding="utf-8", poolclass=NullPool)
        return engine.connect()
    elif dbtype == "MSSQL or " or dbtype == "SQLSERVER":
        engine = create_engine('mssql+pymssql://' + user + ':' + password + '@' + host + ':' + str(port) + '/' + dbname)
        return engine.connect()
    elif dbtype == "INFORMIX":
        engine = create_engine('informix://' + user + ':' + password + '@' + host + ':' + str(port) + '/' + dbname,
                               encoding="utf-8", poolclass=NullPool)
        return engine.connect()
    elif dbtype == "SYBASE":
        engine = create_engine('sybase://' + user + ':' + password + '@' + host + ':' + str(port) + '/' + dbname,
                               encoding="utf-8", poolclass=NullPool)
        return engine.connect()
    elif dbtype == "SQLITE":
        engine = create_engine('ibm_db_sa://' + user + ':' + password + '@' + host + ':' + str(port) + '/' + dbname,
                               encoding="utf-8", poolclass=NullPool)
        return engine.connect()
    elif dbtype == "STS":
        engine = getConnectionSTS(host, port, dbname, user, password)
        return engine
    else:
        return None
        # return conn


def closeConnection(conn):
    if (conn != None):
        conn.close()


'''
查询
conn: 数据库连接
sql: 查询语句
return pd.DataFrame
'''


def query(conn, sql):
    return pd.read_sql_query(sql, con=conn)


'''
将dataFrame插入到数据库
conn: 数据库连接
tableName: 目标表
dataFrame: 要插入的数据框
'''


def dbWriteTable(conn, tableName, dataFrame):
    if isinstance(dataFrame, pd.DataFrame) == False:
        return
    dataFrame.to_sql(tableName, index=False, if_exists='replace', con=conn)


'''
将dataFrame保存到本地
fileName with path: 目标文件名
dataFrame: 要保存的数据框
'''


def dbWriteFileCsv(fileName, dataFrame):
    if isinstance(fileName, pd.DataFrame) == False:
        return
    dataFrame.to_csv(fileName, index=False, header=False)


'''
修改类型
conn:数据库连接
tableName:目标表
dataFrame:要修改的表
Type:修改类型为
'''


def modifyType(conn, tableName, dataFrame, type_dict):
    if isinstance(dataFrame, pd.DataFrame) == False:
        return
    dataFrame.to_sql(tableName, index=False, if_exists='replace', con=conn, dtype=type_dict)


def get_df_from_db(conn,sql):
    cursor = conn.cursor()
    cursor.execute(sql)
    data = cursor.fetchall()
    columnDes = cursor.description #获取连接对象的描述信息
    columnNames = [columnDes[i][0] for i in range(len(columnDes))]
    df = pd.DataFrame([list(i) for i in data],columns=columnNames)
    return df



